package binary_tree.Heap.leetcode;

import java.util.*;

/**
 * @Description: 查找和最小的 K 对数字
 * @Author Stark
 * @Date 2022/3/6 19:55
 **/
public class Num373_SmallestKNumPairs {
    //定义类存储数对
    private class Pairs implements Comparable<Pairs>{
        private int num1;
        private int num2;

        public Pairs(int num1, int num2) {
            this.num1 = num1;
            this.num2 = num2;
        }

        @Override
        public int compareTo(Pairs o) {
            return (o.num1 + o.num2)+  - (this.num1 + this.num2);
        }
    }
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int min1 = Math.min(nums1.length,k);
        int min2 = Math.min(nums2.length,k);
        //使用最大堆存储较小的数对
        Queue<Pairs> queue = new PriorityQueue<>();
        for (int i = 0; i < min1; i++) {
            for (int j = 0; j < min2; j++) {
                if(queue.size() < k){
                    queue.offer(new Pairs(nums1[i],nums2[j]));
                }else{
                    Pairs tmp = queue.peek();
                    int sum = tmp.num1 + tmp.num2;
                    if (nums1[i] + nums2[j] < sum) {
                        queue.poll();
                        queue.offer(new Pairs(nums1[i],nums2[j]));
                    }
                }
            }
        }
        List<List<Integer>> ret = new ArrayList<>();
        for (int i = 0; i < k && (!queue.isEmpty()); i++) {
            Pairs tmp = queue.poll();
            List<Integer> list = new ArrayList<>();
            list.add(tmp.num1);
            list.add(tmp.num2);
            ret.add(list);
        }
        return ret;
    }
}
